{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6e31dbfa",
   "metadata": {},
   "source": [
    "## Financial Research Assistant with Persistent Memory\n",
    "\n",
    "Multi-agent MCPs that researches companies and stores knowledge in a db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdbcb970",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from datetime import datetime\n",
    "from dotenv import load_dotenv\n",
    "from agents.mcp import MCPServerStdio\n",
    "from contextlib import AsyncExitStack\n",
    "from agents import Agent, Runner, trace, Tool\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "036a64c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "research_server_path = os.path.abspath(\"research_server.py\")\n",
    "research_params = {\"command\": \"python\", \"args\": [research_server_path]}\n",
    "\n",
    "brave_env = {\"BRAVE_API_KEY\": os.getenv(\"BRAVE_API_KEY\")}\n",
    "brave_params = {\"command\": \"npx\", \"args\": [\"-y\", \"@modelcontextprotocol/server-brave-search\"], \"env\": brave_env}\n",
    "\n",
    "sandbox_path = os.path.abspath(\"sandbox\")\n",
    "files_params = {\"command\": \"npx\", \"args\": [\"-y\", \"@modelcontextprotocol/server-filesystem\", sandbox_path]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1860c44f",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_researcher(mcp_servers) -> Agent:\n",
    "    instructions = f\"\"\"You research companies using web search tools.\n",
    "Your job is to gather facts about the company: business model, products, recent news, financial metrics.\n",
    "After researching, you must save your findings using the save_company and save_research tools.\n",
    "Always save the company information first with symbol, name, and sector.\n",
    "Then save detailed research findings.\n",
    "Current date: {datetime.now().strftime(\"%Y-%m-%d\")}\n",
    "\"\"\"\n",
    "    return Agent(\n",
    "        name=\"Researcher\",\n",
    "        instructions=instructions,\n",
    "        model=\"gpt-4.1-mini\",\n",
    "        mcp_servers=mcp_servers\n",
    "    )\n",
    "\n",
    "async def get_researcher_tool(mcp_servers) -> Tool:\n",
    "    researcher = await get_researcher(mcp_servers)\n",
    "    return researcher.as_tool(\n",
    "        tool_name=\"Researcher\",\n",
    "        tool_description=\"Research companies and save findings. Provide company symbol and what to research.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702d8a21",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_analyst(mcp_servers) -> Agent:\n",
    "    instructions = \"\"\"You analyze companies and generate investment theses.\n",
    "First, read existing research using the research://company/{symbol} resource.\n",
    "Synthesize the research into a clear investment thesis covering strengths, risks, and outlook.\n",
    "Save your thesis using the save_thesis tool.\n",
    "Create a markdown report and save it to a file.\n",
    "\"\"\"\n",
    "    return Agent(\n",
    "        name=\"Analyst\",\n",
    "        instructions=instructions,\n",
    "        model=\"gpt-4.1-mini\",\n",
    "        mcp_servers=mcp_servers\n",
    "    )\n",
    "\n",
    "async def get_analyst_tool(mcp_servers) -> Tool:\n",
    "    analyst = await get_analyst(mcp_servers)\n",
    "    return analyst.as_tool(\n",
    "        tool_name=\"Analyst\",\n",
    "        tool_description=\"Analyze company research and generate investment thesis. Provide company symbol.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2eb060",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def run_research_assistant(query: str):\n",
    "    async with AsyncExitStack() as stack:\n",
    "        research_server = await stack.enter_async_context(\n",
    "            MCPServerStdio(research_params, client_session_timeout_seconds=60)\n",
    "        )\n",
    "        \n",
    "        brave_server = await stack.enter_async_context(\n",
    "            MCPServerStdio(brave_params, client_session_timeout_seconds=60)\n",
    "        )\n",
    "        \n",
    "        files_server = await stack.enter_async_context(\n",
    "            MCPServerStdio(files_params, client_session_timeout_seconds=60)\n",
    "        )\n",
    "        \n",
    "        researcher_tool = await get_researcher_tool([research_server, brave_server])\n",
    "        analyst_tool = await get_analyst_tool([research_server, files_server])\n",
    "        \n",
    "        main_instructions = \"\"\"You coordinate research and analysis of companies.\n",
    "When asked to research a company:\n",
    "1. Use the Researcher tool to gather and save information\n",
    "2. Use the Analyst tool to create an investment thesis\n",
    "3. Report back with a summary\n",
    "\n",
    "If the company already has research in the database, mention that and still run fresh analysis.\n",
    "\"\"\"\n",
    "        \n",
    "        main_agent = Agent(\n",
    "            name=\"ResearchCoordinator\",\n",
    "            instructions=main_instructions,\n",
    "            model=\"gpt-4.1-mini\",\n",
    "            tools=[researcher_tool, analyst_tool]\n",
    "        )\n",
    "        \n",
    "        with trace(\"research_assistant\"):\n",
    "            result = await Runner.run(main_agent, query, max_turns=30)\n",
    "            display(Markdown(result.final_output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44566f40",
   "metadata": {},
   "source": [
    "### Research a company for the first time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94f397b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "await run_research_assistant(\"Research NVIDIA (NVDA) and provide an investment analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94fc875f",
   "metadata": {},
   "source": [
    "### Check what we have stored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "245c95e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "async with MCPServerStdio(research_params, client_session_timeout_seconds=60) as server:\n",
    "    result = await server.session.read_resource(\"research://company/NVDA\")\n",
    "    for content in result.contents:\n",
    "        print(content.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4036b8b",
   "metadata": {},
   "source": [
    "### Research another company"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45e33f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "await run_research_assistant(\"Research Apple (AAPL) and analyze their business\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c0b022",
   "metadata": {},
   "source": [
    "### View all stored companies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a47c30f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "async with MCPServerStdio(research_params, client_session_timeout_seconds=60) as server:\n",
    "    result = await server.session.read_resource(\"research://all\")\n",
    "    for content in result.contents:\n",
    "        data = json.loads(content.text)\n",
    "        print(f\"Total companies: {data['total_companies']}\")\n",
    "        \n",
    "        for company in data['companies']:\n",
    "            print(f\"- {company['symbol']}: {company['name']} ({company['sector']})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eff774e2",
   "metadata": {},
   "source": [
    "### Re-analyze NVIDIA showing memory persistence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b56fcbe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "await run_research_assistant(\"Update the analysis for NVIDIA with latest information\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
